/*
    Plensa - A Stochastic Simulation Library 

    Copyright (C) 2011  James M. McCollum (james.m.mccollum@gmail.com)

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

package plensa;

import java.util.Vector;
import javax.script.*;

public class Plensa {

    private class Species {
	String name;
	int population;

	public Species(String name, int population) {
	    this.name = name;
	    this.population = population;
	}
    }

    private class Reaction {
	final static int COEFFICIENT = 0;
	final static int SPECIES = 1;
	String equation;
	double rate;
	int reactants[][];
	int products[][];
	double propensity;

	public Reaction(String equation, 
			double rate, 
			int reactants[][], 
			int products[][]) {
	    this.equation = equation;
	    this.rate = rate;
	    this.reactants = reactants;
	    this.products = products;
	}

	public void calculatePropensity() throws Exception {
	    propensity = rate;
	    for(int i=0; i<reactants.length; i++) {
		for(int j=0; j<reactants[i][COEFFICIENT]; j++) {
		    propensity *= speciesList.get(reactants[i][SPECIES]).population - j;
		    propensity /= j+1;
		}
	    }
	}

	public void execute() throws Exception {
	    for(int i=0; i<reactants.length; i++) { 
		speciesList.get(reactants[i][SPECIES]).population -= reactants[i][COEFFICIENT];
		if (speciesList.get(reactants[i][SPECIES]).population < 0) {
		    throw new Exception("ERROR: species population is less than zero."); 
		}
	    }
	    for(int i=0; i<products.length; i++) 
		speciesList.get(products[i][SPECIES]).population += products[i][COEFFICIENT];
	}
    }

    private class CustomReaction extends Reaction {
	String rateEquation;

	ScriptEngineManager factory = new ScriptEngineManager();
	ScriptEngine engine = factory.getEngineByName("JavaScript");

	public CustomReaction(String equation, 
			      String rateEquation, 
			      int reactants[][], 
			      int products[][]) {
	    super(equation,0.0,reactants,products);
	    this.rateEquation = rateEquation;
	}

	public void calculatePropensity() throws Exception {
	    for(Species s : speciesList) 
		engine.put(s.name,s.population);
	    try {
		engine.eval(rateEquation);
	    } catch (Exception e) {
		throw new Exception("Error evaluating rate for equation: " + equation);
	    }
	    try {
		this.propensity = Double.parseDouble(engine.get("rate").toString());
	    } catch (Exception e) {
		throw new Exception("Unable to retrieve value of rate for equation: " + equation);
	    }
	}
    }

    private Vector<Species>  speciesList  = new Vector<Species>();
    private Vector<Reaction> reactionList = new Vector<Reaction>();
    private java.util.Random rand = new java.util.Random();
    private double currentTime = 0.0;
    private double simulatorTime = 0.0;
    private int mu = -1;

    public Plensa() {
	
    }

    public void seed(int value) throws Exception {
	if (currentTime != 0.0) 
	    throw new Exception("unable to reset seed after the simulation has started.");
	rand = new java.util.Random(value);
    }

    public void species(String name, int population) throws Exception {
	if (currentTime != 0.0) 
	    throw new Exception("unable to add species after the simulation has started.");
	if (!isValidName(name)) 
	    throw new Exception("invalid species name: " + name);
	if (findSpecies(name) != -1) 
	    throw new Exception("duplicate species name: " + name);
	if (population < 0) 
	    throw new Exception("negative population for species " + name + ": " + population);

	speciesList.add(new Species(name,population));
    }

    public void reaction(String equation, double rate) throws Exception {
	if (currentTime != 0.0) 
	    throw new Exception("unable to add reactions after the simulation has started.");
	if (rate <= 0.0) 
	    throw new Exception("invalid rate for equation '" + equation + "': " + rate); 
	String parts[] = equation.split("=>");
	if (parts.length != 2) 
	    throw new Exception("unable to find '=>' in equation '" + equation + "'");
	int reactants[][] = parseSpeciesList(equation,parts[0]);
	int products[][] = parseSpeciesList(equation,parts[1]);
	reactionList.add(new Reaction(equation,rate,reactants,products));
    }

    public void reaction(String equation, String rateEquation) throws Exception {
	if (currentTime != 0.0) 
	    throw new Exception("unable to add reactions after the simulation has started.");
	String parts[] = equation.split("=>");
	if (parts.length != 2) 
	    throw new Exception("unable to find '=>' in equation '" + equation + "'");
	int reactants[][] = parseSpeciesList(equation,parts[0]);
	int products[][] = parseSpeciesList(equation,parts[1]);
	reactionList.add(new CustomReaction(equation,rateEquation,reactants,products));
    }

    public void simulate(double dt) throws Exception {
	if (dt < 0)
	    throw new Exception("called simulate with a change in time less than 0.");

	while (simulatorTime < currentTime+dt) {
	    if (mu != -1) {
		reactionList.get(mu).execute();
	    }

	    double totalPropensity = 0.0;
	    for(Reaction reaction : reactionList) {
		reaction.calculatePropensity();
		totalPropensity += reaction.propensity;
	    }

	    double selector = rand.nextDouble()*totalPropensity;
	    for(mu=0; mu<reactionList.size(); mu++) {
		selector -= reactionList.get(mu).propensity;
		if (selector <= 0.0) break;
	    }
	    if (mu == reactionList.size())
		throw new Exception("ERROR: mu == reactionList.size()");
	    
	    simulatorTime += -Math.log(rand.nextDouble()) / totalPropensity;
	}

	currentTime += dt;
    }

    public void header() {
	System.out.print("time");
	for(Species s : speciesList) 
	    System.out.print("\t" + s.name);
	System.out.println();
    }

    public void output() {
	System.out.print(time());
	for(Species s : speciesList) 
	    System.out.print("\t" + s.population);
	System.out.println();
    }

    public double time() {
	return currentTime;
    }

    public int population(String name) throws Exception {
	int species = findSpecies(name);
	if (species == -1) throw new Exception("invalid species: " + name);
	return speciesList.get(species).population;
    }

    private int findSpecies(String name) {
	for(int i=0; i<speciesList.size(); i++) {
	    if (speciesList.get(i).name.equals(name)) {
		return i;
	    }
	}
	return -1;
    }

    private boolean isValidName(String name) {
	if (name.length() == 0) return false;
	if (!Character.isLetter(name.charAt(0))) return false;
	for(int i=1; i<name.length(); i++) {
	    if ((!Character.isLetterOrDigit(name.charAt(i))) && (name.charAt(i) != '_')) return false;
	}
	return true;
    }

    private int[][] parseSpeciesList(String equation, String str) throws Exception {
	str = str.trim();
	if (str.equals("*")) {
	    int x[][] = new int[0][2];
	    return x;
	}

	String strs[] = str.split("\\+");
	int coefficients[] = new int[speciesList.size()];
	int count = 0;
	for(int i=0; i<strs.length; i++) {
	    int speciesIndex = findSpecies(strs[i].trim());
	    if (speciesIndex == -1) 
		    throw new Exception("unknown species '" + strs[i].trim() + "' in equation '" + equation + "'");
	    if (coefficients[speciesIndex] == 0) count++;
	    coefficients[speciesIndex]++;
	}

	int x[][] = new int[count][2];
	int j=0;
	for(int i=0; i<speciesList.size(); i++) {
	    if (coefficients[i] != 0) {
		x[j][Reaction.COEFFICIENT] = coefficients[i];
		x[j][Reaction.SPECIES] = i;
		j++;
	    }
	}

	return x;
    }

    public static void main(String args[]) {
	new Plensa();
    }
}

